{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Playwright Target - optional\n",
    "\n",
    "This notebook demonstrates how to interact with the **Playwright Target** in PyRIT.\n",
    "\n",
    "The `PlaywrightTarget` class allows you to interact with web applications using\n",
    "[Playwright](https://playwright.dev/python/docs/intro).\n",
    "This is useful for testing interactions with web applications, such as chatbots or other web interfaces,\n",
    "especially for red teaming purposes.\n",
    "\n",
    "## Example Setup\n",
    "\n",
    "Before you begin, ensure you have the correct version of PyRIT installed and any necessary secrets configured as\n",
    "described [here](../../setup/populating_secrets.md).\n",
    "\n",
    "To run the Flask app, you also must [download](https://ollama.com/download) and run Ollama, making sure the flask is using a correct model. For example, `ollama run llama3.2:1b` runs the Llama 3.2 1B model.\n",
    "\n",
    "Additionally, you need to install playwright by executing `playwright install`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "## Example: Interacting with a Web Application using `PlaywrightTarget`\n",
    "\n",
    "In this example, we'll interact with a simple web application running locally at `http://127.0.0.1:5000`.\n",
    "which runs a chatbot that responds to user prompts via ollama\n",
    "We'll define an interaction function that navigates to the web application, inputs a prompt, and retrieves the\n",
    "bot's response.\n",
    "\n",
    "## Start the Flask App\n",
    "Before we can interact with the web application, we need to start the Flask app that serves the chatbot, this will be done in a subprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Flask app from /workspace/doc/code/targets/playwright_demo/app.py...\n",
      "Flask app is running and ready!\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import sys\n",
    "import time\n",
    "from urllib.error import URLError\n",
    "from urllib.request import urlopen\n",
    "\n",
    "\n",
    "def start_flask_app() -> subprocess.Popen:\n",
    "    # Get the notebook's directory\n",
    "    notebook_dir = os.getcwd()\n",
    "\n",
    "    # Construct the path to app.py relative to the notebook directory\n",
    "    app_path = os.path.join(notebook_dir, \"playwright_demo\", \"app.py\")\n",
    "\n",
    "    # Ensure that app.py exists\n",
    "    if not os.path.isfile(app_path):\n",
    "        raise FileNotFoundError(f\"Cannot find app.py at {app_path}\")\n",
    "\n",
    "        # Start the Flask app\n",
    "    print(f\"Starting Flask app from {app_path}...\")\n",
    "    flask_process = subprocess.Popen([sys.executable, app_path], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\n",
    "\n",
    "    # Wait for the app to start by checking if the server is up\n",
    "    server_url = \"http://127.0.0.1:5000\"\n",
    "    server_ready = False\n",
    "    while not server_ready:\n",
    "        try:\n",
    "            response = urlopen(server_url)\n",
    "            if response.status == 200:\n",
    "                server_ready = True\n",
    "        except URLError:\n",
    "            time.sleep(0.5)  # Wait a bit before checking again\n",
    "    print(\"Flask app is running and ready!\")\n",
    "    return flask_process\n",
    "\n",
    "\n",
    "# Start the Flask app\n",
    "flask_process = start_flask_app()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "The flask app should now be running locally:\n",
    "\n",
    "![image-2.png](../../../assets/playwright_demo.png)\n",
    "\n",
    "### Interaction Function\n",
    "This is playwright script that interacts with the chatbot web application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from playwright.async_api import Page, async_playwright\n",
    "\n",
    "from pyrit.models import Message\n",
    "from pyrit.prompt_target import PlaywrightTarget\n",
    "from pyrit.setup import IN_MEMORY, initialize_pyrit\n",
    "\n",
    "initialize_pyrit(memory_db_type=IN_MEMORY)\n",
    "\n",
    "\n",
    "# Define the interaction function\n",
    "async def interact_with_my_app(page: Page, message: Message) -> str:\n",
    "    # Define selectors\n",
    "    input_selector = \"#message-input\"\n",
    "    send_button_selector = \"#send-button\"\n",
    "    bot_message_selector = \".bot-message\"\n",
    "\n",
    "    # Count existing messages\n",
    "    initial_messages = await page.query_selector_all(bot_message_selector)\n",
    "    initial_message_count = len(initial_messages)\n",
    "\n",
    "    # Wait for the page to be ready\n",
    "    await page.wait_for_selector(input_selector)\n",
    "\n",
    "    # Send the prompt text (get from first request piece)\n",
    "    prompt_text = message.message_pieces[0].converted_value\n",
    "    await page.fill(input_selector, prompt_text)\n",
    "    await page.click(send_button_selector)\n",
    "\n",
    "    # Wait for new messages (bot messages)\n",
    "    await page.wait_for_function(\n",
    "        f\"document.querySelectorAll('{bot_message_selector}').length > {initial_message_count}\"\n",
    "    )\n",
    "\n",
    "    # Extract the bot's response text\n",
    "    bot_message_element = await page.query_selector(f\"{bot_message_selector}:last-child\")\n",
    "    bot_response = await bot_message_element.text_content()\n",
    "    return bot_response.strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {
    "lines_to_next_cell": 0
   },
   "source": [
    "### Using `PlaywrightTarget` with the Interaction Function and Scorer\n",
    "\n",
    "Now, we can use the `PlaywrightTarget` by passing the interaction function we defined.\n",
    "We'll use the `PromptSendingAttack` to send prompts to the target and collects responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[34m🔹 Turn 1 - USER\u001b[0m\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[34m  Tell me a joke about computer programming.\u001b[0m\n",
      "\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[1m\u001b[33m🔸 ASSISTANT\u001b[0m\n",
      "\u001b[33m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n",
      "\u001b[33m  Howdy there, partner! I've got one for ya that's fixin' to boot up your funny bone. Here goes:\u001b[0m\n",
      "\u001b[33m  \u001b[0m\n",
      "\u001b[33m    Why did the programmer quit his job?\u001b[0m\n",
      "\u001b[33m  \u001b[0m\n",
      "\u001b[33m    Because he didn't get arrays enough pay!\u001b[0m\n",
      "\u001b[33m  \u001b[0m\n",
      "\u001b[33m    I know, I know, it's a little corny, but I'm a cowboy who code, what can I say?\u001b[0m\n",
      "\n",
      "\u001b[34m────────────────────────────────────────────────────────────────────────────────────────────────────\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import sys\n",
    "\n",
    "from pyrit.executor.attack import ConsoleAttackResultPrinter, PromptSendingAttack\n",
    "\n",
    "if sys.platform == \"win32\":\n",
    "    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())\n",
    "\n",
    "\n",
    "# Using PlaywrightTarget with the interaction function and scorer\n",
    "async def main(page: Page) -> None:\n",
    "    target = PlaywrightTarget(interaction_func=interact_with_my_app, page=page)\n",
    "\n",
    "    attack = PromptSendingAttack(objective_target=target)\n",
    "\n",
    "    objective = \"Tell me a joke about computer programming.\"\n",
    "\n",
    "    result = await attack.execute_async(objective=objective)  # type: ignore\n",
    "    await ConsoleAttackResultPrinter().print_conversation_async(result=result)  # type: ignore\n",
    "\n",
    "\n",
    "async def run() -> None:\n",
    "    async with async_playwright() as playwright:\n",
    "        browser = await playwright.chromium.launch(headless=True)\n",
    "        context = await browser.new_context()\n",
    "        page: Page = await context.new_page()\n",
    "        await page.goto(\"http://127.0.0.1:5000\")\n",
    "        await main(page)\n",
    "        await context.close()\n",
    "        await browser.close()\n",
    "\n",
    "\n",
    "# Note in Windows this doesn't run in jupyter notebooks due to playwright limitations\n",
    "# https://github.com/microsoft/playwright-python/issues/480\n",
    "await run()\n",
    "\n",
    "# if __name__ == \"__main__\":\n",
    "#     asyncio.run(run())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Terminate the Flask App"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Flask app has been terminated.\n"
     ]
    }
   ],
   "source": [
    "# Terminate the Flask app when done\n",
    "flask_process.terminate()\n",
    "flask_process.wait()  # Ensure the process has terminated\n",
    "print(\"Flask app has been terminated.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Close connection to memory\n",
    "from pyrit.memory import CentralMemory\n",
    "\n",
    "memory = CentralMemory.get_memory_instance()\n",
    "memory.dispose_engine()"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
